#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文档翻译模块
包含所有与文档翻译相关的功能
"""

import os
import re
import yaml
import json
import tempfile
import time
import requests
from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH, WD_LINE_SPACING
from docx.oxml.ns import qn
from docx.oxml import OxmlElement

# 配置加载函数
def load_config():
    """从config.yaml加载配置信息"""
    config_path = os.path.join(os.path.dirname(__file__), "config.yaml")
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        return config
    except Exception as e:
        print(f"加载配置文件失败: {e}")
        # 返回默认配置
        return {
            "model_config": {
                "api_url": "http://10.2.7.150:8000/v1/chat/completions",
                "api_key": "sk-4f9d5fe0a9644f11af996c40de98af77",
                "model_name": "qwen3-vl-8b",
                "max_tokens": 2048,
                "stream": True
            },
            "translation_config": {
                "prompt_template": "将以下英文准确翻译成中文，保持专业术语的准确性，不添加任何解释或额外内容：\n{text}",
                "delay_between_requests": 0.1
            }
        }

# 加载配置
config = load_config()

# 从配置中获取模型信息
QWEN_API_URL = config["model_config"]["api_url"]
API_KEY = config["model_config"]["api_key"]
MODEL_NAME = config["model_config"]["model_name"]
MAX_TOKENS = config["model_config"]["max_tokens"]
STREAM_ENABLED = config["model_config"]["stream"]
PROMPT_TEMPLATES = config["translation_config"]["prompt_templates"]
DEFAULT_DIRECTION = config["translation_config"]["default_direction"]
REQUEST_DELAY = config["translation_config"]["delay_between_requests"]

def translate_text(text, translation_direction=None):
    """使用Qwen3VL API翻译文本（使用流式处理，确保完整接收）"""
    if not text.strip():
        return text
    
    # 确定翻译方向
    if translation_direction is None:
        translation_direction = DEFAULT_DIRECTION
    
    # 获取对应的prompt模板
    prompt_template = PROMPT_TEMPLATES.get(translation_direction, PROMPT_TEMPLATES[DEFAULT_DIRECTION])
        
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    # 构造请求参数：启用流式返回
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt_template.format(text=text)
                    }
                ]
            }
        ],
        "max_tokens": MAX_TOKENS,  # 从配置中获取token上限
        "stream": STREAM_ENABLED  # 从配置中获取流式设置
    }
    
    try:
        # 发送流式请求
        response = requests.post(QWEN_API_URL, json=payload, headers=headers, stream=True)
        response.raise_for_status()  # 检查HTTP请求是否成功
        
        # 收集流式响应内容
        full_response = ""
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')
                # 检查是否为数据行（以"data: "开头）
                if line_str.startswith('data: '):
                    data_content = line_str[6:]  # 去掉"data: "前缀
                    
                    # 检查是否为结束标记（[DONE]）
                    if data_content.strip() == '[DONE]':
                        break
                    
                    try:
                        # 解析JSON数据
                        data_json = json.loads(data_content)
                        
                        # 检查是否有choices字段和delta字段
                        if "choices" in data_json and len(data_json["choices"]) > 0:
                            choice = data_json["choices"][0]
                            if "delta" in choice and "content" in choice["delta"]:
                                content = choice["delta"]["content"]
                                if content:
                                    full_response += content
                    except json.JSONDecodeError:
                        # 忽略JSON解析错误，继续处理下一行
                        continue
        
        # 如果成功获取到翻译结果，返回翻译后的文本
        if full_response.strip():
            print(f"流式API响应: {full_response}")  # 打印流式响应内容，方便调试
            time.sleep(REQUEST_DELAY)  # 从配置中获取延迟时间
            return full_response.strip()
        else:
            # 如果流式处理失败，回退到非流式处理
            print("流式处理未返回有效内容，尝试非流式处理")
            return translate_text_non_streaming(text)
            
    except Exception as e:
        print(f"流式翻译出错: {str(e)}")
        print(f"API响应内容: {response.text if 'response' in locals() else '无响应'}")  # 打印响应便于调试
        # 出错时回退到非流式处理
        return translate_text_non_streaming(text)

def translate_text_non_streaming(text, translation_direction=None):
    """非流式翻译函数（作为流式处理的备用方案）"""
    if not text.strip():
        return text
    
    # 确定翻译方向
    if translation_direction is None:
        translation_direction = DEFAULT_DIRECTION
    
    # 获取对应的prompt模板
    prompt_template = PROMPT_TEMPLATES.get(translation_direction, PROMPT_TEMPLATES[DEFAULT_DIRECTION])
        
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt_template.format(text=text)
                    }
                ]
            }
        ],
        "max_tokens": MAX_TOKENS,
        "stream": False
    }
    
    try:
        response = requests.post(QWEN_API_URL, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()
        
        if "choices" in result and len(result["choices"]) > 0:
            translated_content = result["choices"][0]["message"]["content"]
            if isinstance(translated_content, list) and len(translated_content) > 0:
                return translated_content[0].get("text", text).strip()
            elif isinstance(translated_content, str):
                return translated_content.strip()
        
        return text
    except Exception as e:
        print(f"非流式翻译出错: {str(e)}")
        return text

def process_paragraph(paragraph, translate_tables=False, translation_direction=None):
    """处理段落，翻译文本同时保留格式和图片"""
    # 保存原始格式信息
    original_style = paragraph.style
    original_align = paragraph.alignment
    
    # 检查段落中是否包含图片
    has_images = False
    
    # 使用更可靠的图片检测方法
    for run in paragraph.runs:
        # 检查run中是否包含图片元素
        if (run._element.xpath('.//pic:pic') or 
            run._element.xpath('.//w:drawing') or 
            run._element.xpath('.//a:blip')):
            has_images = True
            break
    
    # 如果段落包含图片，采用更保守的处理方式：只翻译文本，不修改图片部分
    if has_images:
        print(f"检测到段落包含图片，采用保守处理方式")
        
        # 提取纯文本内容（不包括图片）
        text_runs = []
        for run in paragraph.runs:
            if run.text and run.text.strip():
                text_runs.append(run)
        
        # 如果有文本需要翻译
        if text_runs:
            # 收集所有文本内容
            original_text = ' '.join([run.text for run in text_runs if run.text.strip()])
            if original_text.strip():
                translated_text = translate_text(original_text, translation_direction)
            else:
                translated_text = ""
            
            # 只替换文本run的内容，保留图片run不变
            if translated_text:
                # 清除原有文本run的内容
                for run in text_runs:
                    run.text = ""
                
                # 在第一个文本run中添加翻译后的文本
                if text_runs:
                    text_runs[0].text = translated_text
    else:
        # 如果不包含图片，使用改进的处理方式
        original_text = paragraph.text
        if not original_text.strip():
            return  # 空段落不处理
            
        translated_text = translate_text(original_text, translation_direction)
        
        # 改进的格式保留方法：保留原有run结构，只替换文本内容
        if paragraph.runs:
            # 如果有多个run，合并到第一个run中
            first_run = paragraph.runs[0]
            
            # 保存第一个run的格式信息
            font_name = first_run.font.name
            font_size = first_run.font.size
            font_bold = first_run.font.bold
            font_italic = first_run.font.italic
            font_underline = first_run.font.underline
            font_color = first_run.font.color.rgb
            
            # 清除所有run的内容
            for run in paragraph.runs:
                run.text = ""
            
            # 只在第一个run中添加翻译后的文本
            first_run.text = translated_text
            
            # 恢复第一个run的格式
            first_run.font.name = font_name
            first_run.font.size = font_size
            first_run.font.bold = font_bold
            first_run.font.italic = font_italic
            first_run.font.underline = font_underline
            if font_color:
                first_run.font.color.rgb = font_color
            
            # 设置中文字体
            if 'font' in first_run._element:
                rPr = first_run._element.font
                rFonts = OxmlElement('w:rFonts')
                rFonts.set(qn('w:eastAsia'), 'SimSun')  # 设置宋体为中文字体
                rPr.insert_element_before(rFonts, 'w:b', 'w:i', 'w:color', 'w:sz', 'w:u')
        else:
            # 如果没有run，创建新的run
            new_run = paragraph.add_run(translated_text)
            
            # 应用原始样式
            paragraph.style = original_style
            paragraph.alignment = original_align
            
            # 设置基本字体格式
            new_run.font.name = 'Times New Roman'
            new_run.font.size = Pt(12)
            
            # 设置中文字体
            if 'font' in new_run._element:
                rPr = new_run._element.font
                rFonts = OxmlElement('w:rFonts')
                rFonts.set(qn('w:eastAsia'), 'SimSun')  # 设置宋体为中文字体
                rPr.insert_element_before(rFonts, 'w:b', 'w:i', 'w:color', 'w:sz', 'w:u')
    
    # 设置段落格式以允许西文在中间换行
    paragraph_format = paragraph.paragraph_format
    
    # 保留原始对齐方式，不强制设置为左对齐
    # paragraph.alignment = WD_ALIGN_PARAGRAPH.LEFT
    
    # 设置允许西文在单词中间换行
    # 通过设置wordWrap属性为0来允许单词内换行
    pPr = paragraph._element.pPr
    if pPr is None:
        pPr = OxmlElement('w:pPr')
        paragraph._element.insert(0, pPr)
    
    # 添加wordWrap元素
    wordWrap = OxmlElement('w:wordWrap')
    wordWrap.set(qn('w:val'), '0')  # 0表示允许单词内换行
    pPr.append(wordWrap)
    
    # 设置合理的行间距
    paragraph_format.line_spacing_rule = WD_LINE_SPACING.SINGLE
    paragraph_format.space_before = Pt(0)
    paragraph_format.space_after = Pt(6)
    
    # 禁用孤行控制
    paragraph_format.widow_control = False

def process_table(table, translation_direction=None):
    """处理表格，翻译文本同时保留格式和图片"""
    for row in table.rows:
        for cell in row.cells:
            for paragraph in cell.paragraphs:
                process_paragraph(paragraph, translation_direction=translation_direction)

def get_document_paragraph_count(input_path, translate_tables=False):
    """获取文档中的总工作量（段落数 + 表格单元格数）"""
    doc = Document(input_path)
    total_work = len(doc.paragraphs)
    
    # 如果用户选择翻译表格，则包含表格单元格数量
    if translate_tables:
        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    # 每个单元格至少包含一个段落
                    total_work += len(cell.paragraphs)
    
    return total_work

def translate_document(input_path, output_path, translate_tables=False, translation_direction=None, progress_callback=None):
    """翻译整个文档，保留排版和图片，支持进度回调"""
    doc = Document(input_path)
    
    # 获取总工作量（段落数 + 表格单元格数）
    total_work = get_document_paragraph_count(input_path, translate_tables)
    completed_work = 0
    
    # 处理段落
    for i, paragraph in enumerate(doc.paragraphs):
        # 在每次处理前检查是否应该中断
        if progress_callback:
            # 通过回调函数检查是否应该中断
            try:
                progress = int(completed_work / total_work * 100)
                progress_callback(progress, completed_work, total_work)
            except Exception as e:
                # 如果回调函数抛出异常（表示任务被取消），则中断翻译
                if "翻译任务已被用户取消" in str(e):
                    raise e
                else:
                    # 其他异常继续处理
                    pass
        
        process_paragraph(paragraph, translate_tables=translate_tables, translation_direction=translation_direction)
        completed_work += 1
        
        # 更新进度
        if progress_callback:
            progress = int(completed_work / total_work * 100)
            progress_callback(progress, completed_work, total_work)
    
    # 处理表格（仅在用户选择翻译表格时执行）
    if translate_tables:
        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    # 处理单元格中的每个段落
                    for paragraph in cell.paragraphs:
                        # 在每次处理前检查是否应该中断
                        if progress_callback:
                            try:
                                progress = int(completed_work / total_work * 100)
                                progress_callback(progress, completed_work, total_work)
                            except Exception as e:
                                if "翻译任务已被用户取消" in str(e):
                                    raise e
                                else:
                                    pass
                        
                        process_paragraph(paragraph, translate_tables=translate_tables, translation_direction=translation_direction)
                        completed_work += 1
                        
                        # 更新进度
                        if progress_callback:
                            progress = int(completed_work / total_work * 100)
                            progress_callback(progress, completed_work, total_work)
    
    # 保存翻译后的文档
    doc.save(output_path)
    
    return output_path

def extract_images_from_docx(docx_path, output_dir="extract"):
    """从DOCX文档中提取图片到指定目录"""
    import zipfile
    import shutil
    import os
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 解压DOCX文件（DOCX本质上是ZIP格式）
    with zipfile.ZipFile(docx_path, 'r') as zip_ref:
        zip_ref.extractall(output_dir)
    
    # 查找图片文件（通常在word/media目录中）
    media_dir = os.path.join(output_dir, "word", "media")
    if os.path.exists(media_dir):
        image_files = [f for f in os.listdir(media_dir) 
                      if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff'))]
        
        # 将图片复制到extract目录的根目录
        for img_file in image_files:
            src_path = os.path.join(media_dir, img_file)
            dst_path = os.path.join(output_dir, img_file)
            shutil.copy2(src_path, dst_path)
        
        # 清理解压的临时文件
        shutil.rmtree(os.path.join(output_dir, "word"))
        
        return image_files
    
    return []

def get_document_images_info(docx_path):
    """获取文档中图片的信息和位置"""
    doc = Document(docx_path)
    
    image_info = []
    
    # 检查段落中的图片
    for i, paragraph in enumerate(doc.paragraphs):
        for run in paragraph.runs:
            if run._element.xpath('.//pic:pic'):
                image_info.append({
                    'type': 'paragraph',
                    'position': i + 1,
                    'text': paragraph.text[:100] + '...' if len(paragraph.text) > 100 else paragraph.text
                })
    
    # 检查表格中的图片
    for table_idx, table in enumerate(doc.tables):
        for row_idx, row in enumerate(table.rows):
            for cell_idx, cell in enumerate(row.cells):
                for para in cell.paragraphs:
                    for run in para.runs:
                        if run._element.xpath('.//pic:pic'):
                            image_info.append({
                                'type': 'table',
                                'position': f"表格{table_idx+1}-行{row_idx+1}-列{cell_idx+1}",
                                'text': para.text[:100] + '...' if len(para.text) > 100 else para.text
                            })
    
    return image_info